from node import *


def wide_and_deep(n, classes=2, hiddens=(12,)):
    """
    构造一个Wide & Deep模型的计算图，特征数量为 n 。 n 个特征能产生
    C(n, 2) 个二次交互项 xi * xj 。
    """

    # x 是一个 n 维向量变量，不初始化，不参与训练
    x = Variable((n, 1), init=False, trainable=False)

    # 二次特征，共C(n, 2)=n * (n-1) / 2个
    assert n > 1
    power2_n = int(n * (n - 1) / 2)
    x_2 = Variable((power2_n, 1), init=False, trainable=False)

    # Wide部分
    wide = Add(
        MatMul(
            Variable((classes, power2_n), True),
            x_2
        ),
        Variable((classes, 1), True)
    )

    # Deep部分
    input_size = n
    input = x
    for h_size in hiddens:
        output = ReLU(
            Add(
                MatMul(
                    Variable((h_size, input_size), True),
                    input),
                Variable((h_size, 1), True)
            )
        )
        input_size = h_size
        input = output

    # 输出层的神经元
    deep = Add(
        MatMul(
            Variable((classes, input_size), True),
            output),
        Variable((classes, 1), True)
    )

    # 将Wide部分和Deep部分的输出相加，得到logits
    logits = Add(wide, deep)

    # 返回输入和logits
    return x, x_2, logits
